import React, { useRef, useState, useCallback } from 'react';
import { produce } from 'immer';
import type { ChatItem, MessageEnd } from '../types';
import { XStream } from '@ant-design/x';
import MyFetch from '../fetch.js';

function unicodeToChar(text: string) {
  if (!text)
    return ''

  return text.replace(/\\u[0-9a-f]{4}/g, (_match, p1) => {
    return String.fromCharCode(parseInt(p1, 16))
  })
}

const useSSE = (props: any) => {
  const { prevChatList, } = props;
  const conversationId = useRef('')
  const hasStopResponded = useRef(false)
  const taskIdRef = useRef('')
  const isRespondingRef = useRef(false)
  const [lines, setLines] = React.useState<Record<string, string>[]>([])
  const [isResponding, setIsResponding] = useState(false)
  const [chatList, setChatList] = useState<ChatItem[]>(prevChatList || [])
  const chatListRef = useRef<ChatItem[]>(prevChatList || [])

  const handleUpdateChatList = useCallback((newChatList: ChatItem[]) => {
    setChatList(newChatList)
    chatListRef.current = newChatList
  }, [])

  const handleResponding = useCallback((isResponding: boolean) => {
    setIsResponding(isResponding)
    isRespondingRef.current = isResponding
  }, [])

  const updateCurrentQA = useCallback(({
    responseItem,
    questionId,
    placeholderAnswerId,
    questionItem,
  }: {
    responseItem: ChatItem
    questionId: string
    placeholderAnswerId: string
    questionItem: ChatItem
  }) => {
    const newListWithAnswer = produce(
      chatListRef.current.filter(item => item.id !== responseItem.id && item.id !== placeholderAnswerId),
      (draft) => {
        if (!draft.find((item: any) => item.id === questionId))
          draft.push({ ...questionItem })

        draft.push({ ...responseItem })
      })
    handleUpdateChatList(newListWithAnswer)
  }, [handleUpdateChatList])

  const ssePost = (data: any) => {
    const questionId = `question-${Date.now()}`
    const questionItem = {
      id: questionId,
      content: data.query,
      isAnswer: false,
      message_files: data.files,
    }

    const placeholderAnswerId = `answer-placeholder-${Date.now()}`
    const placeholderAnswerItem = {
      id: placeholderAnswerId,
      content: '',
      isAnswer: true,
    }
    const newList = [...chatListRef.current, questionItem, placeholderAnswerItem]
    handleUpdateChatList(newList)

    // answer
    const responseItem: ChatItem = {
      id: placeholderAnswerId,
      content: '',
      agent_thoughts: [],
      message_files: [],
      isAnswer: true,
    }

    handleResponding(true)
    hasStopResponded.current = false
    let hasSetResponseId = false

    const handler = {
      onData: (message: string, isFirstMessage: boolean, { conversationId: newConversationId, messageId, taskId }: any) => {
        responseItem.content = responseItem.content + message

        if (messageId && !hasSetResponseId) {
          responseItem.id = messageId
          hasSetResponseId = true
        }

        if (isFirstMessage && newConversationId)
          conversationId.current = newConversationId

        taskIdRef.current = taskId
        if (messageId)
          responseItem.id = messageId

        updateCurrentQA({
          responseItem,
          questionId,
          placeholderAnswerId,
          questionItem,
        })
      },

      onMessageEnd() {
        const newListWithAnswer = produce(
          chatListRef.current.filter(item => item.id !== responseItem.id && item.id !== placeholderAnswerId),
          (draft) => {
            if (!draft.find(item => item.id === questionId))
              draft.push({ ...questionItem })

            draft.push({ ...responseItem })
          })
        handleUpdateChatList(newListWithAnswer)
      },

      async onCompleted(hasError?: boolean) {
        handleResponding(false)

        if (hasError)
          return
      },
    };
    request(data.query, handler)
  }

  const request = async (msg: string, options?: any) => {
    const response = await MyFetch({
      query: msg,
    });
    // .....
  
    const res = XStream({
      readableStream: response.body,
    })
    handleStream(res, options);
  }

  const handleStream = async (res: any, options?: any) => {
    const { onData, onMessageEnd, onCompleted } = options
    let isFirstMessage = true;
    let bufferObj: Record<string, any> = {};
    for await (const chunk of res) {
      setLines((pre) => [...pre, chunk]);
      try {
        if (chunk.data) {
          bufferObj = JSON.parse(chunk.data) as Record<string, any>
        }
      } catch (err) {
        console.log(err, '出错了')
        onData('', isFirstMessage, {
          conversationId: bufferObj?.conversation_id,
          messageId: bufferObj?.message_id,
        })
        return
      }
      if (bufferObj.event === 'message' || bufferObj.event === 'agent_message') {
        // can not use format here. Because message is splitted.
        onData(unicodeToChar(bufferObj.answer), isFirstMessage, {
          conversationId: bufferObj.conversation_id,
          taskId: bufferObj.task_id,
          messageId: bufferObj.id,
        })
        isFirstMessage = false
      }
      else if (bufferObj.event === 'message_end') {
        onMessageEnd?.(bufferObj as MessageEnd)
      }
    }
    onCompleted();
  }

  return {
    lines,
    chatList,
    isResponding,
    ssePost,
  }
};

export default useSSE;
